#include "checksum.h"
#include <assert.h>
#include <stdint.h>
#include <stdlib.h>
#include <cstdio>
#include <string.h>

uint32_t getChecksum32(uint32_t base, uint8_t *addr, int len) {
    // packet's checksum is 0
    uint32_t sum = base;
    while (len > 1) {
        uint16_t v = *(uint16_t *) addr;
        // printf("v: %x\n", htons(v));
        sum += htons(v);
        len -= 2;
        addr += 2;
    }
    if (len > 0) {
        sum += *addr << 8;
        //printf("v: %x\n", *addr);
    }
    return sum;
}

uint16_t getChecksum16(uint32_t checksum) {
    while (checksum > 0xffff) {
        checksum = (checksum & 0xffff) + (checksum >> 16);
    }
    return checksum;
}

bool validateAndFillChecksum(uint8_t *packet, size_t len) {
    struct ip6_hdr *ip6 = (struct ip6_hdr *) packet;
    // check next header
    uint16_t nxt_header = ip6->ip6_nxt;
    uint32_t checksum = 0;
    checksum = getChecksum32(checksum, (uint8_t * ) & ip6->ip6_src, 16);
    checksum = getChecksum32(checksum, (uint8_t * ) & ip6->ip6_dst, 16);
    checksum = getChecksum32(checksum, (uint8_t * ) & ip6->ip6_plen, 2);
    checksum += nxt_header;

    bool res = false;
    if (nxt_header == IPPROTO_UDP) {
        // UDP
        struct udphdr *udp = (struct udphdr *) &packet[sizeof(struct ip6_hdr)];
        // length: udp->uh_ulen
        // checksum: udp->uh_sum
        uint16_t packet_checksum = ntohs(udp->uh_sum);
        udp->uh_sum = 0;
        checksum = getChecksum32(checksum, (uint8_t *) udp, ntohs(ip6->ip6_plen));
        checksum = getChecksum16(checksum);
        checksum = 0xffff - checksum;

        udp->uh_sum = htons(checksum);
        if (checksum == 0) { // 对于发送方，如果计算出的校验和为 0，则需要设置校验和字段为 0xFFFF
            udp->uh_sum = 0xffff;
        }

        if (packet_checksum == 0) {
            // 如果packet_checksum为0，意味着没有进行校验
            // 可以直接返回false
            return false;
        }
        res = checksum == packet_checksum;
        if (checksum == 0 && packet_checksum == 0xffff) {
            // 特别处理计算出的校验和为0的场景
            res = true;
        }
        if (checksum == 0xffff) {
            // 这种情形会出现吗？实验文档说不会...
            assert(0);
        }
    } else if (nxt_header == IPPROTO_ICMPV6) {
        // ICMPv6
        struct icmp6_hdr *icmp =
                (struct icmp6_hdr *) &packet[sizeof(struct ip6_hdr)];
        // length: len-sizeof(struct ip6_hdr)
        // checksum: icmp->icmp6_cksum
        uint16_t packet_checksum = ntohs(icmp->icmp6_cksum);
        icmp->icmp6_cksum = 0;
        checksum = getChecksum32(checksum, (uint8_t *) icmp, ntohs(ip6->ip6_plen));
        checksum = getChecksum16(checksum);
        checksum = 0xffff - checksum;
        if (checksum == 0 && packet_checksum == 0xffff) {
            icmp->icmp6_cksum = 0x0;
            return true;
        }
        if (checksum == packet_checksum) {
            res = true;
        }
        icmp->icmp6_cksum = htons(checksum);
    } else {
        printf("next_header: %d\n", nxt_header);
        assert(false);
    }
    return res;
}
